import numpy as np
import matplotlib
import matplotlib.pyplot as plt


# 从csv文件中，读取储存的数据
def load_csvdata(filename, x_scale=1.0, y_scale=1.0, z_scale=1.0):
    '''
    :param filename:  必填项，字符串类型，文件名称
    :param x_scale:   选填项，浮点型，放大比例
    :param y_scale:   选填项，浮点型，放大比例
    :param z_scale:   选填项，浮点型，放大比例
    :return:  返回处理后的横坐标，纵坐标，竖坐标的矩阵，（一维，一维，二维）
    '''
    data = np.loadtxt(filename, delimiter=',')  # 读取数据
    x = data[0, 1:]   # 读取横坐标
    y = data[1:, 0]   # 读取纵坐标
    z = data[1:, 1:]  # 读取数据场
    return x * x_scale, y * y_scale, z * z_scale  # 返回数据放大后的数据


# 空间坐标处理(将离散的步长处理成为离散的距离，返回距离值和最大距离)
def step_sum(grid_step):
    grid_sum = grid_step.cumsum()  # 累计求和
    grid_half = 0.5 * grid_step  # 网格取半
    gird_middle = grid_sum - grid_half  # 取中间值
    return gird_middle  # 返回中间距离值


# 画等高线图表
def plot_contour(x, y, z,
                 field_style=None,
                 transposition: bool = False,
                 **kwargs):
    figure = plt.figure()  # 创建图片
    ax = figure.add_subplot(111)  # 添加图表1

    # 横纵坐标是否转置
    if transposition:
        x, y = y, x
        z = z.T  # 场转置

    x, y = np.meshgrid(x, y)

    # 根据场选择标记颜色
    if field_style == 'rh':
        levels = [10, 45, 50, 55, 60, 65, 70, 75, 80, 100]
        colormap_style = 'PuBu'
    elif field_style == 't':
        levels = [0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50]
        colormap_style = 'OrRd'
    elif field_style == 'p':
        levels = 10
        colormap_style = 'bwr'
    else:
        levels = 5
        colormap_style = 'rainbow'

    if kwargs == {}:
        pass
    else:
        levels = kwargs['levels']
        colormap_style = kwargs['cmp']
    image = ax.contourf(x, y, z, levels, alpha=0.8, cmap=colormap_style,)

    if field_style == 'rh':  # 湿度场划线
        image2 = ax.contour(image, levels=[0,80,100], colors='#2b2b2b', linewidths=0.2)
        # ax.clabel(image2, inline=True, fontsize=8, fmt='%d%%')
    plt.colorbar(image)  # 对image图片设置增加彩色图例
    return ax  # 返回画布,图画


# 设置轴标题
def set_axis_fontlabel(ax, xlabel, ylabel, title=None):
    font_ttc = 'C:\Windows\Fonts\simsun.ttc'  # 读取宋体字体路径
    font = matplotlib.font_manager.FontProperties(fname=font_ttc)  # 设置字体，用于显示汉字
    # 设置坐标轴名称，图表名称
    if title is not None:
        ax.set_title(title, fontproperties=font)
    ax.set_xlabel(xlabel, fontproperties=font)
    ax.set_ylabel(ylabel, fontproperties=font)


# 程序中的输出图片
def plot_field_contour(filename, field_style='rainbow',
                       layers_thickness: list or np.ndarray = None,
                       labletext: list = ['时间/d', '距离/mm'], **kwargs):

    '''
    :param filename: 必填项，保存计算程序输出数据的文件名
    :param field_style: 选填项，颜色色彩选择，默认‘rainbow’，
                        还有：‘rh’湿度场，‘t’温度场，
                        和其它colormap的样式选择
    :param layers_thickness: 选填项，列表[数值]格式，界面分界线
    :param labletext: 选填项，列表[横坐标标题，纵坐标标题，图表标题（可不填）]
    :return: 返回对象，画布Axes
    '''
    x_scale = 1000
    y_scale = 1/3600/24
    z_scale = 1
    x, y, z = load_csvdata(filename, x_scale, y_scale, z_scale)  # 读取数据并放大
    x = step_sum(x)  # 自身特殊网格处理，将离散的步长变为距离
    transposition = True  # 是否转置xy轴
    ax = plot_contour(x, y, z, field_style, transposition, **kwargs)  # 画等高线图
    # 绘制界面线
    if layers_thickness is not None:
        layers_thickness = np.array(layers_thickness).cumsum()
        for i in layers_thickness:
            yy = [y[0], y[-1]]
            xx = [i*x_scale, i*x_scale]
            if transposition:
                xx, yy = yy, xx
            ax.plot(xx, yy, '--', linewidth=1, color='k', alpha=0.3)

    set_axis_fontlabel(ax, *labletext)  # 设置轴标题
    ax.invert_yaxis()  # y轴坐标倒置
    return ax


if __name__ == '__main__':
    import os
    from tkinter.filedialog import askopenfilename
    # 读取文件路径
    filepath = askopenfilename(filetypes=[("csv file", "*.csv")])
    foldername, filename = os.path.split(filepath)  # 分割文件夹与文件名
    fieldname, extension = os.path.splitext(filename)  # 分割文件名与后缀

    field_style = 'rh'
    layers_thickness = [0.01, 0.08, 0.24]
    levels = [10, 45, 50, 55, 60, 65, 70, 75, 80,85,90, 100]
    ax = plot_field_contour(filepath, field_style, layers_thickness,
                            levels=levels, cmp='PuBu'
                            )

    plt.savefig('{}/contour_{}.png'.format(foldername, fieldname), dpi=350)

    plt.show()  # 显示图片